import os

import pytorch_lightning.utilities.seed
import torch
import torch.nn as nn
from PIL import Image
from pytorch_lightning import LightningModule, Trainer
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as f
from torchvision import transforms


class TestModule(LightningModule):
    def __init__(self):
        super(TestModule, self).__init__()
        self.loss_function = nn.CrossEntropyLoss()

    def forward(self, x):
        pass

    def training_step(self, batch, batch_index):
        pass

    def configure_optimizers(self):
        pass


class MyData(Dataset):
    def __init__(self):
        super().__init__()

    def __len__(self):
        pass

    def __getitem__(self, index):
        pass


model = TestModule()
model = model.to("cuda")
dataset = MyData()
trainer = Trainer(accelerator="auto", gpus=1, max_epochs=100)
dataLoader = DataLoader(dataset, batch_size=32)
trainer.fit(model, train_dataloaders=dataLoader)



